%% On the Optimal Design of Transfers and Income-Tax Progressivity
% Numerical evaluation of the analytical welfare formula

clear; clc; close all;
addpath('../Matlab_Figure_Files');

%% Parameters and Grids 

% Parameters
phi = 1/0.4;            % inverse Frisch (standard value)
tau = 0.181;            % benchmark progressivity (from HSV 2017)
n_0_target = 0.3;       % target for labor supply without transfers
B = n_0_target^(-1-phi)*(1-tau);    % labor disutility parameter to match target for labor supply
G = 0.233*n_0_target;   % spending (targeted spending output ratio 0.233)
rhoa = 0.935;           % persistence idiosyncratic shock
v_log_con = 0.18;                           % variance log consumption data (from HSV 2017)
vw   = (1-rhoa^2)*v_log_con*(1/(1-tau)^2);  % variance idiosyncratic shock to match variance of log consumption in the data
vw_aux = vw/(1-rhoa^2);                     % translation from AR(1) to lognormal shocks
lambda_calib = (n_0_target-G)/(n_0_target)*exp(tau*(1-tau)*vw/((1-rhoa^2)));  % lambda in the calibration (rescaled)

disp('Calibration')
disp(['Implied vw to match variance of log consumption is = ',num2str(vw)])
disp(['Expressed for lognormal shocks rather than AR(1) = ',num2str(vw_aux)])
disp(['Calibrated lambda is = ',num2str(lambda_calib)])
disp(' ')

% Grids
tau_vec = -0.9:0.01:0.9;                    % grid for progressivity
n_0 = ((1-tau_vec)./B).^(1/(1+phi));        % labor supply depending on progressivity (no transfer)
vc = ((1-tau_vec).*vw)./(2*(1-rhoa^2));     % auxiliary variable

%% Welfare without Transfers

% Welfare components without transfers
eff_nt = log(n_0-G);                % efficiency term without transfers
ls_nt = -(1-tau_vec)./(1+phi);      % labor supply term without transfers
rd_nt = -(1-tau_vec).*vc;           % redistribution term without transfers

% Welfare without transfers
W_nt_ha = eff_nt + ls_nt + rd_nt;   % all terms
[~,aux_ha] = max(W_nt_ha); gopt_nt_ha   = tau_vec(aux_ha);
W_nt_ra = eff_nt + ls_nt;           % representative agent
[~,aux_ra] = max(W_nt_ra); gopt_nt_ra   = tau_vec(aux_ra);
W_nt_ra_nog = log(n_0) + ls_nt;     % representative agent and zero spending
[~,aux_ra_nog] = max(W_nt_ra_nog); gopt_nt_ra_nog   = tau_vec(aux_ra_nog);
disp('Welfare without transfers: Optimal progressivity for different cases')
disp(['Optimal progressivity tau is                  = ',num2str(gopt_nt_ha)])
disp(['Optimal progressivity tau with RA is          = ',num2str(gopt_nt_ra)])
disp(['Optimal progressivity tau with RA and no G is = ',num2str(gopt_nt_ra_nog)])
disp(' ')

%% Welfare with Transfers: Representative Agent, First Best

% For calibrated G
T_vec = linspace(-0.10,0.10,101);               % grid for transfer
T_vec_norm = T_vec./n_0_target;                 % normalized by calibrated mean income
n_foc   = @(n) B*n^phi - 1/(n-G);               % labor first best condition 
disp('RA first best: Solve nonlinear equation for labor supply')
n_fb    = fsolve(@(n)n_foc(n),1);               % labor first best
disp(' ')
tau_opt_fb_ra = -(G+T_vec)./(n_fb-(G+T_vec));   % first best tau as function of T

% For higher G
G_aux = G*1.5;                                  % higher G
n_foc   = @(n) B*n^phi - 1/(n-G_aux);           % labor first best condition 
disp('RA first best with higher G: Solve nonlinear equation for labor supply')
disp(['Higher G is ',num2str(G_aux/G),' times baseline G'])
n_fb_highG    = fsolve(@(n)n_foc(n),1);         % labor first best
disp(' ')
tau_opt_fb_ra_highG = -(G_aux+T_vec)./(n_fb_highG-(G_aux+T_vec));   % first best tau as function of T

% Plot for paper - color
fig = figure(101);
plot(T_vec_norm*100,tau_opt_fb_ra,'LineWidth',2); hold on
plot(T_vec_norm*100,tau_opt_fb_ra_highG,'-.','Color',[0.5,0.5,0.5],'LineWidth',2); hold on
plot([0;0], [-0.8; gopt_nt_ra],'--','Color',[0.8500    0.3250    0.0980],'LineWidth',2); hold on
plot([-G/n_0_target*100;-G/n_0_target*100], [-0.8; 0], ':','Color',[0.9290, 0.6940, 0.1250],'LineWidth',2); hold on
plot([-0.3333333333*100;-G/n_0_target*100], [0;0], ':','Color',[0.9290, 0.6940, 0.1250],'LineWidth',2); hold on
plot([-0.3333333333*100;0], [-0.26;gopt_nt_ra],'--','Color',[0.8500    0.3250    0.0980],'LineWidth',2); hold on
leg = legend('Calibrated spending $G$','Higher spending $\hat{G}$','Location','NorthEast');
set(leg,'Interpreter','LaTex','Fontsize',12);
legend boxoff
xlim([-0.3333333333*100 0.33333333333*100])
ylim([-0.8 0.2])
set(gca,'XGrid','off','YGrid','on','Fontsize',12)  
set(gca,'TickLabelInterpreter','LaTex')
xlabel('$T$ in percent of calibrated mean income','Interpreter','LaTex','Fontsize',12)
ylabel('$\tau$','Interpreter','LaTex','Fontsize',12)
run graph_extended.m;
saveas(gcf,'am_ra_fb.pdf')

% Plot for paper - bw
fig = figure(102);
plot(T_vec_norm*100,tau_opt_fb_ra,'-','Color',[0 0 0],'LineWidth',2); hold on
plot(T_vec_norm*100,tau_opt_fb_ra_highG,'-.','Color',[0.5,0.5,0.5],'LineWidth',2); hold on
plot([0;0], [-0.8; gopt_nt_ra],':','Color',[0 0 0],'LineWidth',2); hold on
plot([-G/n_0_target*100;-G/n_0_target*100], [-0.8; 0], '--','Color',[0 0 0],'LineWidth',2); hold on
plot([-0.3333333333*100;-G/n_0_target*100], [0;0], '--','Color',[0 0 0],'LineWidth',2); hold on
plot([-0.3333333333*100;0], [-0.26;gopt_nt_ra],':','Color',[0 0 0],'LineWidth',2); hold on
leg = legend('Calibrated spending $G$','Higher spending $\hat{G}$','Location','NorthEast');
set(leg,'Interpreter','LaTex','Fontsize',12);
legend boxoff
xlim([-0.3333333333*100 0.33333333333*100])
ylim([-0.8 0.2])
set(gca,'XGrid','off','YGrid','on','Fontsize',12)  
set(gca,'TickLabelInterpreter','LaTex')
xlabel('$T$ in percent of calibrated mean income','Interpreter','LaTex','Fontsize',12)
ylabel('$\tau$','Interpreter','LaTex','Fontsize',12)
run graph_extended.m;
saveas(gcf,'am_ra_fb_bw.pdf')
saveas(gcf,'am_ra_fb_bw','epsc')

%% Welfare with Transfers: Approximation for Heterogeneous Agents

% Additional welfare components with transfers
T = 0.01;   % small transfer
Omega_e_ra = -(n_0./(n_0-G)).*(1/(1+phi)).*(1./(n_0-G)) + ((1-tau_vec)./(1+phi)).*(1./(n_0-G)); % efficiency term RA with transfers
Omega_e_ha = ((tau_vec.*(1-tau_vec))./(n_0-G)).*(vw/(1-rhoa^2)).*(n_0./(n_0-G)).*(1/(1+phi));   % efficiency term HA with transfers
Omega_r = (((1-tau_vec).^2)./(n_0-G)).*(vw/(1-rhoa^2));                                         % redistribution term with transfers

% Welfare with transfers
W_wt_ha = W_nt_ha + T*(Omega_e_ra+Omega_e_ha+Omega_r);      % all terms
[~,aux_ha_wt] = max(W_wt_ha); gopt_wt_ha   = tau_vec(aux_ha_wt);
W_wt_ra = W_nt_ra + T*Omega_e_ra;                           % representative agent
[~,aux_ra_wt] = max(W_wt_ra); gopt_wt_ra   = tau_vec(aux_ra_wt);
disp('Welfare with transfers: Optimal progressivity for different cases')
disp(['Optimal progressivity tau with RA and T = 0 is = ',num2str(gopt_nt_ra)])
disp(['Optimal progressivity tau with RA and T > 0 is = ',num2str(gopt_wt_ra)])
disp(['Optimal progressivity tau with HA and T = 0 is = ',num2str(gopt_nt_ha)])
disp(['Optimal progressivity tau with HA and T > 0 is = ',num2str(gopt_wt_ha)])
disp(' ')

% Bound tau_hat(G) described in paper
tau_hat_G = 1 - B*G^(1+phi)*((2*(1+phi))/(2*(1+phi)-1))^(1+phi);
disp(['Bound tau_hat(G) is = ', num2str(tau_hat_G)])